package main

import (
	"encoding/csv"
	"fmt"
	"os"
	"strconv"
)

// 常量定义
const (
	InitialCapital   = 100000.0 // 初始资金10万元
	BuyRatio         = 0.25     // 买入比例25%
	AddPositionRatio = 0.25     // 加仓比例25%
	LossThreshold    = -0.2     // 加仓触发亏损率30%
	ProfitThreshold  = 0.1      // 止盈比例10%
	CommissionRate   = 0.001    // 手续费率0.1%
	MinShares        = 100      // 最小交易单位100股
	MA5Period        = 5        // MA5周期
	MA20Period       = 20       // MA20周期
)

// 股票数据结构
type StockData struct {
	Date   string
	Open   float64
	High   float64
	Low    float64
	Close  float64
	Change float64
	MA5    float64
	MA20   float64
}

// 持仓信息
type Position struct {
	Shares      int     // 持仓数量
	Cost        float64 // 持仓成本
	TotalCost   float64 // 累计投入成本
	MarketValue float64 // 当前市值
	Profit      float64 // 持仓盈亏
	ProfitRatio float64 // 持仓盈亏比
	TradeCount  int     // 交易次数
	HoldDays    int     // 持股天数
}

// 账户信息
type Account struct {
	Cash             float64  // 可用资金
	TotalAsset       float64  // 总资产(现金+持仓市值)
	Position         Position // 持仓信息
	LastBuyDate      string   // 最后买入日期
	LastBuyDateIndex int      // 最后买入日期索引
}

func main() {
	// 1. 读取CSV数据
	data, err := readCSV("F:\\stock\\每日涨停\\588080.csv")
	if err != nil {
		fmt.Printf("读取CSV失败: %v\n", err)
		return
	}

	// 2. 计算均线
	calculateMA(data)

	// 3. 执行回测
	results := backtest(data)

	// 4. 输出结果
	err = writeResults(results, "trade.txt")
	if err != nil {
		fmt.Printf("输出结果失败: %v\n", err)
		return
	}

	fmt.Println("回测完成，结果已输出到trade.txt")
}

// 读取CSV文件
func readCSV(filePath string) ([]StockData, error) {
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	reader := csv.NewReader(file)
	records, err := reader.ReadAll()
	if err != nil {
		return nil, err
	}

	// 跳过表头
	records = records[1:]

	var data []StockData
	for _, record := range records {
		if len(record) < 6 {
			continue
		}

		open, _ := strconv.ParseFloat(record[1], 64)
		high, _ := strconv.ParseFloat(record[2], 64)
		low, _ := strconv.ParseFloat(record[3], 64)
		close, _ := strconv.ParseFloat(record[4], 64)
		change, _ := strconv.ParseFloat(record[5], 64)

		data = append(data, StockData{
			Date:   record[0],
			Open:   open,
			High:   high,
			Low:    low,
			Close:  close,
			Change: change,
		})
	}

	return data, nil
}

// 计算均线
func calculateMA(data []StockData) {
	for i := range data {
		// 计算MA5
		if i >= MA5Period-1 {
			sum := 0.0
			for j := i - MA5Period + 1; j <= i; j++ {
				sum += data[j].Close
			}
			data[i].MA5 = sum / float64(MA5Period)
		}

		// 计算MA20
		if i >= MA20Period-1 {
			sum := 0.0
			for j := i - MA20Period + 1; j <= i; j++ {
				sum += data[j].Close
			}
			data[i].MA20 = sum / float64(MA20Period)
		}
	}
}

// 执行回测
func backtest(data []StockData) []string {
	account := Account{
		Cash:       InitialCapital,
		TotalAsset: InitialCapital,
	}

	var results []string
	results = append(results, "日期\t开盘价\t最高价\t最低价\t收盘价\t涨幅\t总资产\t总市值\t成本价\t总盈亏金额\t总盈亏比\t持仓盈亏金额\t持仓盈亏比\t交易次数\t持股天数")

	for i, day := range data {
		// 更新持仓市值和持股天数
		if account.Position.Shares > 0 {
			account.Position.MarketValue = float64(account.Position.Shares) * day.Close
			account.Position.Profit = account.Position.MarketValue - account.Position.TotalCost
			account.Position.ProfitRatio = account.Position.Profit / account.Position.TotalCost
			account.TotalAsset = account.Cash + account.Position.MarketValue
			account.Position.HoldDays++
		}

		// 最后交易日强制平仓
		if i == len(data)-1 && account.Position.Shares > 0 {
			sellAll(&account, day)
		}

		// 检查卖出条件 (持仓盈利≥10%且不是买入当天)
		if account.Position.Shares > 0 &&
			account.Position.ProfitRatio >= ProfitThreshold &&
			account.LastBuyDate != day.Date {
			sellAll(&account, day)
		}

		// 检查买入条件
		if i >= MA20Period-1 { // 确保有足够的均线数据
			// 无持仓时买入
			if account.Position.Shares == 0 && day.MA5 > day.MA20 {
				buy(&account, day)
			}

			// 有持仓时加仓
			if account.Position.Shares > 0 &&
				account.Position.ProfitRatio <= LossThreshold &&
				account.Cash > 0 {
				buy(&account, day)
			}
		}

		// 记录当日结果
		results = append(results, formatDailyResult(day, account))
	}

	return results
}

// 买入操作
func buy(account *Account, day StockData) {
	// 计算可买入金额
	buyAmount := account.TotalAsset * BuyRatio
	if buyAmount > account.Cash {
		buyAmount = account.Cash
	}

	// 计算可买入股数 (100股整数倍)
	shares := int(buyAmount/day.Close/MinShares) * MinShares
	if shares < MinShares {
		return // 不足100股不交易
	}

	// 计算手续费
	commission := float64(shares) * day.Close * CommissionRate
	totalCost := float64(shares)*day.Close + commission

	// 更新账户
	account.Cash -= totalCost
	account.Position.Shares += shares
	account.Position.TotalCost += float64(shares) * day.Close
	account.Position.Cost = account.Position.TotalCost / float64(account.Position.Shares)
	account.Position.TradeCount++
	account.LastBuyDate = day.Date
}

// 卖出操作
func sellAll(account *Account, day StockData) {
	if account.Position.Shares == 0 {
		return
	}

	// 计算卖出金额
	amount := float64(account.Position.Shares) * day.Close
	commission := amount * CommissionRate
	netAmount := amount - commission

	// 更新账户
	account.Cash += netAmount
	account.TotalAsset = account.Cash
	account.Position.TradeCount++ // 增加交易次数
	account.Position = Position{} // 清空持仓
}

// 格式化每日结果
func formatDailyResult(day StockData, account Account) string {
	holdDays := 0
	if account.Position.Shares > 0 {
		holdDays = account.Position.HoldDays
	}
	return fmt.Sprintf("%s\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f%%\t%.2f\t%.2f\t%.2f\t%.2f\t%.2f%%\t%.2f\t%.2f%%\t%d\t%d",
		day.Date,
		day.Open,
		day.High,
		day.Low,
		day.Close,
		day.Change*100,
		account.TotalAsset,
		account.Position.MarketValue,
		account.Position.Cost,
		account.TotalAsset-InitialCapital,
		(account.TotalAsset/InitialCapital-1)*100,
		account.Position.Profit,
		account.Position.ProfitRatio*100,
		account.Position.TradeCount,
		holdDays,
	)
}

// 输出结果到文件
func writeResults(results []string, filePath string) error {
	file, err := os.Create(filePath)
	if err != nil {
		return err
	}
	defer file.Close()

	for _, line := range results {
		_, err := file.WriteString(line + "\n")
		if err != nil {
			return err
		}
	}

	return nil
}
